package main

import (
    "log"
    "os"
    "io"
    "errors"
    "strings"
    "io/ioutil"
    "crypto/cipher"
    "crypto/hmac"
    "encoding/base64"
    "crypto/rand"
    "crypto/rsa"
    "crypto/aes"
    "crypto/sha256"
    "crypto/x509"
    "encoding/asn1"
    "encoding/pem"
    mrand "math/rand"
)

type Crypt struct {
    pubKey       string             /* Public key stirng   */
    privateKey   string             /* Private key string  */
    key         *rsa.PrivateKey     /* Private key         */
}

func createKey() Crypt {
    reader := rand.Reader
    bitSize := 2048

    key, err := rsa.GenerateKey(reader, bitSize)
    if err != nil {
        panic(err)
    }

    /* Save public key */
    asn1Bytes, err := asn1.Marshal(key.PublicKey)
    if err != nil {
        panic(err)
    }
    mem := pem.EncodeToMemory(&pem.Block{
	Type:  "PUBLIC KEY",
	Bytes: asn1Bytes,
    })
    log.Printf("%s", mem)

    /* Save private key */
    mem2 := pem.EncodeToMemory(&pem.Block{
        Type: "RSA PRIVATE KEY",
        Bytes: x509.MarshalPKCS1PrivateKey(key),
    })

    return Crypt{
        pubKey: string(mem),
        privateKey: string(mem2),
        key: key,
    }
}

func restoreKey() Crypt {
    var keyPath string = "1.pem"
    var key2Path string = "2.pem"
    var saveKey bool = true
    /* Step 1. Restore RSA key  */
    content, err1 := ioutil.ReadFile(keyPath)
    content2, err2 := ioutil.ReadFile(key2Path)
    if err1 == nil && err2 == nil {

        block, err3 := pem.Decode([]byte(content2))
        if block == nil {
            panic(err3)
        }

        priv, err4 := x509.ParsePKCS1PrivateKey(block.Bytes)
        if err4 != nil {
            panic(err4)
        }

        return Crypt{
            pubKey: string(content),
            privateKey: string(content2),
            key: priv,
        }
    }
    /* Step 2. Generate RSA key */
    key := createKey()
    /* Step 2. Save RSA key */
    if saveKey {
        /* Save pulic key */
        pemfile, err := os.Create(keyPath)
        if err != nil {
            panic(err)
        }
        pemfile.Write([]byte(key.pubKey))
        defer pemfile.Close()
        /* Save pulic key */
        pemfile2, err2 := os.Create(key2Path)
        if err2 != nil {
            panic(err2)
        }
        pemfile2.Write([]byte(key.privateKey))
        defer pemfile2.Close()
    }
    return key
}

func randint(min, max int) int {
    return min + mrand.Intn(max - min)
}

type Crypticle struct {
    aesKey  string       /* AES key     */
    hmacKey string       /* HMAC key    */
}

const SIG_SIZE = 32

func NewCrypticle(keyStr string) *Crypticle {
    
    var keySize int = 192
    
    data, err := base64.StdEncoding.DecodeString(keyStr)
    log.Printf("crypticle key %s %v", data, err)
    
    if len(data) != keySize / 8 + SIG_SIZE {
        panic("invalid key")
    }
    var size int = len(data)
    var aesKey string = string(data[:size-SIG_SIZE])
    var hmacKey string = string(data[size-SIG_SIZE:])

    log.Printf("key1 = %s key2 = %s", aesKey, hmacKey)

    return &Crypticle{
        aesKey: aesKey,
        hmacKey: hmacKey,
    }
}

const AES_BLOCK_SIZE = 16
const PICKLE_PAD = "pickle::"

func (c Crypticle) Encode(data string) (string, error) {

    /* Step 1. Padding input */
//    var pad int = AES_BLOCK_SIZE - len(data) % AES_BLOCK_SIZE
    var pad int = aes.BlockSize - len(data) % aes.BlockSize
    log.Printf("Padding %d", pad)
    var ch_pad string = string(rune(pad))
    log.Printf("Padding by %s", ch_pad)
    log.Printf("Before %s", data)
    data = data + strings.Repeat(ch_pad, pad)
    log.Printf("After %s", data)

    /* Step 2. Encrypt */
    block, err := aes.NewCipher([]byte(c.aesKey))
    if err != nil {
        return "", err
    }
    ciphertext := make([]byte, aes.BlockSize + len(data))
    iv := ciphertext[:aes.BlockSize]
    if _, err := io.ReadFull(rand.Reader, iv); err != nil {
	panic(err)
    }
    mode := cipher.NewCBCEncrypter(block, iv)
    var raw []byte = []byte(data)
    mode.CryptBlocks(ciphertext[aes.BlockSize:], raw)

    /* Step 3. Sign with SHA256 message digest */
    mac := hmac.New(sha256.New, []byte(c.hmacKey))
    mac.Write(ciphertext)
    sig := mac.Sum(nil)
    log.Printf("sig = %q", sig)

    /* Step 4. Wrap result */
    var result string = string(ciphertext) + string(sig)
    log.Printf("result = %q", result)
    return result, nil
}

func (c Crypticle) Decode(msg string) (string, error) {

    var size int = len(msg)

    /* Step 1. Unwrap container with signature */
    var data string = msg[:size-SIG_SIZE]
//    var sig string = msg[size-SIG_SIZE:]

    /* Step 2. Calculate summary */
    size = len(data)
    log.Printf("Block count: %d", size / AES_BLOCK_SIZE)

    /* Step 3. Unwrap start vector */
    var iv_bytes []byte = []byte( data[:AES_BLOCK_SIZE] )
    //data = data[AES_BLOCK_SIZE:]

    ciphertext := []byte( data[AES_BLOCK_SIZE:] )

    block, err := aes.NewCipher([]byte(c.aesKey))
    if err != nil {
        return "", err
    }
//    ciphertext = ciphertext[aes.BlockSize:]

    // CBC mode always works in whole blocks
    if len(ciphertext) % aes.BlockSize != 0 {
        return "", errors.New("ciphertext is not a multiple of the block size")
    }

    mode := cipher.NewCBCDecrypter(block, iv_bytes)
    mode.CryptBlocks(ciphertext, ciphertext)

    //
    s := string(ciphertext[:])
    log.Printf("decode = %s", s)

    if strings.HasPrefix(s, PICKLE_PAD) {
        var prefix_size int = len(PICKLE_PAD)
        new_ret := s[prefix_size:]
        return new_ret, nil
    }

    return "", nil
}

